import cv2
import numpy as np
from ultralytics import YOLO

# Load the YOLO model
model = YOLO("model2/yolo11n-seg.pt")

# Open the video file
video_path = "video/0705/test2/frame_00136.jpg"
cap = cv2.VideoCapture(video_path)

def check_line_mask_intersection(x1, y1, x2, y2, mask):
    """
    检查给定的直线是否与mask相交。
    
    :param x1: 直线起点x坐标
    :param y1: 直线起点y坐标
    :param x2: 直线终点x坐标
    :param y2: 直线终点y坐标
    :param mask: 二值化的人体mask
    :return: 如果直线与mask相交则返回True，否则返回False
    """
    original_h, original_w = mask.shape[:2]
    points_on_line = np.linspace((x1, y1), (x2, y2), num=100, dtype=np.int32)
    for x, y in points_on_line:
        if 0 <= y < original_h and 0 <= x < original_w:
            if mask[y, x] > 0:
                return True
    return False

# Loop through the video frames
while cap.isOpened():
    # Read a frame from the video
    success, frame = cap.read()
    original_h, original_w = frame.shape[:2]
    print(original_h, original_w)

    # 直线起点终点
    x1, y1 = 0,1657
    x2, y2 = 3840,1657


    # 沿线采样
    points = np.linspace((x1, y1), (x2, y2), 100, dtype=np.int32)

    if success:
        # Run YOLO inference on the frame
        results = model(frame,classes=[0],imgsz=1280)

        for result in results:
            boxes = result.boxes.xyxy.cpu().numpy().tolist()
            xy = result.masks.xy  # mask in polygon format
            xyn = result.masks.xyn  # normalized
            masks = result.masks.data.cpu().numpy() # mask in matrix format (num_objects x H x W)

        # print(boxes)
        # print(len(boxes))
        # print(masks)
        # print(len(masks))

        for i, person_mask in enumerate(masks):
            person_mask_resized = cv2.resize(person_mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
            intersects = check_line_mask_intersection(x1, y1, x2, y2, person_mask_resized)
            print(f"人体{i+1} {'与直线相交' if intersects else '不与直线相交'}")
                # Visualize the results on the frame
        annotated_frame = results[0].plot()

        cv2.line(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.imwrite("output.jpg", annotated_frame)

        break
# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()